# -*- coding: utf-8 -*-
# @time: 2024/6/5 16:21
# @file: IV_RV
# @author: tyshixi08
import pandas as pd

from get_data.origin_data import *
from get_data.RSJ import get_bond_return
import os

# 获取转债退市信息字典
def bond_out_market():
    data = convertible.instruments(code_ls())
    dic = {}
    for item in data:
        if item['stop_trading_date'] is not None:
            dic[item['order_book_id']] = item['stop_trading_date'].strftime("%Y-%m-%d")
    return dic

# 获取可转债强赎信息字典
def bond_get_call():
    df = convertible.get_call_info(code_ls(), str(month_ls()[0]).replace('-', '') + '01', str(int(str(month_ls()[-1]).replace('-', ''))) + '01').reset_index()[['order_book_id', 'info_date']]
    df['info_date'] = df['info_date'].astype(str)
    dic = df.set_index('order_book_id')['info_date'].to_dict()
    return dic

# 获取只需要进行退市处理的可转债代码列表
def bond_out_market_filter_ls():
    ls = list(set(bond_out_market().keys()) - set(bond_get_call().keys()))
    return ls

# 获取只需要进行强赎处理的可转债代码列表
def bond_get_call_filter_ls():
    ls = list(set(bond_get_call().keys()) - set(bond_out_market().keys()))
    return ls

# 获取既需要进行退市处理，又需要进行强赎处理的可转债代码列表
def bond_out_market_get_call_filter_ls():
    ls = list(set(bond_out_market()) & set(bond_get_call()))
    return ls

# 获取无需处理的可转债代码列表
def bond_code_ls():
    ls = list(set(code_ls()) - set(bond_out_market_filter_ls()) - set(bond_get_call_filter_ls()) - set(bond_out_market_get_call_filter_ls()))
    return ls

# 根据可转债代码获取IV
def get_IV(order_book_ids, sta_date = month_ls()[0] + '-01', end_date = month_ls()[-1] + '-01', flag = 1):
    df = convertible.get_indicators(order_book_ids = order_book_ids, start_date=sta_date, end_date=end_date,fields=['iv','remaining_size']).reset_index()
    if flag == 1:
        df = df[df['iv'] > 0.01].dropna()
        return df
    elif flag == 0:
        return df

# 退市调整IV
def IV_out_market_filter():
    df = get_IV(bond_out_market_filter_ls())

    # 退市数据调整
    for category in df['order_book_id'].unique():
        lines_keep = df[df['order_book_id'] == category].shape[0] - 22
        df_keep = df[df['order_book_id'] == category].head(lines_keep)
        df = df[df['order_book_id'] != category]
        df = pd.concat([df, df_keep])

    df = df.dropna()
    return df

# 强赎调整IV
def IV_get_call_filter():
    df = get_IV(bond_get_call_filter_ls())

    # 强赎信息调整
    dic_get_call = bond_get_call()
    for category in df['order_book_id'].unique():
        df_keep = df[df['order_book_id'] == category]
        df_keep = df_keep[df_keep['date'] < dic_get_call[category]]
        df = df[df['order_book_id'] != category]
        df = pd.concat([df, df_keep])

    df = df.dropna()
    return df

# 退市与强赎调整IV
def IV_out_market_get_call_filter():
    df = get_IV(bond_out_market_get_call_filter_ls())
    dic_get_call = bond_get_call()

    for category in df['order_book_id'].unique():
        # 退市数据调整
        lines_keep = df[df['order_book_id'] == category].shape[0] - 22
        df_keep = df[df['order_book_id'] == category].head(lines_keep)
        df = df[df['order_book_id'] != category]
        df = pd.concat([df, df_keep])

        # 强赎信息调整
        df_keep = df[df['order_book_id'] == category]
        df_keep = df_keep[df_keep['date'] < dic_get_call[category]]
        df = df[df['order_book_id'] != category]
        df = pd.concat([df, df_keep])

    df = df.dropna()
    return df

# 获取经调整后的所有可转债IV
def get_IV_filter():
    df = pd.concat([get_IV(bond_code_ls()), IV_out_market_filter(), IV_get_call_filter(), IV_out_market_get_call_filter()])
    df = pd.merge(df, get_all_convertible_code(), how = 'outer', on = 'order_book_id').dropna()
    return df.sort_values('date').reset_index(drop='True')[['order_book_id', 'stock_code', 'date', 'iv']]

##############################################################################################################################################################

# 计算可转债正股的日收益率
def get_stock_return(stock_code):
    df = get_bond_return(code=stock_code, start_date = month_ls()[0].replace('-','') + '01', end_date = month_ls()[-1].replace('-','') + '01')
    return df

# 计算可转债正股的实际波动率(滚动5天)
def get_stock_RV_5D(stock_code):
    df = get_stock_return(stock_code)
    df_RV = df.set_index('date').groupby('order_book_id')['return'].rolling(window=5).std() * math.sqrt(252)
    df_RV = df_RV.reset_index().dropna().rename(columns={'return':'RV'})
    df = pd.merge(df, df_RV, how='outer', on=['order_book_id', 'date']).rename(columns={'order_book_id':'stock_code'}).dropna()[['stock_code', 'return', 'date', 'RV']]
    return df.drop_duplicates()

# 计算可转债正股的实际波动率(滚动22天)
def get_stock_RV_22D(stock_code):
    df = get_stock_return(stock_code)
    df_RV = df.set_index('date').groupby('order_book_id')['return'].rolling(window=22).std() * math.sqrt(252)
    df_RV = df_RV.reset_index().dropna().rename(columns={'return':'RV'})
    df = pd.merge(df, df_RV, how='outer', on=['order_book_id', 'date']).rename(columns={'order_book_id':'stock_code'}).dropna()[['stock_code', 'return', 'date', 'RV']]
    return df.drop_duplicates()

# 获取5D隐波差
def IV_RV_5D():
    df_bond = get_IV_filter()
    df_bond['date'] = pd.to_datetime(df_bond['date']).dt.strftime('%Y-%m-%d')
    df_stock = get_stock_RV_5D(list(set(df_bond['stock_code'])))
    df_stock['date'] = pd.to_datetime(df_stock['date']).dt.strftime('%Y-%m-%d')
    df = pd.merge(df_bond, df_stock, how='inner', on=['stock_code', 'date'])
    df = df[['order_book_id', 'stock_code', 'date', 'iv', 'RV']]
    df['IVdelta_mid_5D'] = df['iv'] - df['RV']
    df = df.dropna()
    return df

# 获取22D隐波差
def IV_RV_22D():
    df_bond = get_IV_filter()
    df_bond['date'] = pd.to_datetime(df_bond['date']).dt.strftime('%Y-%m-%d')
    df_stock = get_stock_RV_22D(list(set(df_bond['stock_code'])))
    df_stock['date'] = pd.to_datetime(df_stock['date']).dt.strftime('%Y-%m-%d')
    df = pd.merge(df_bond, df_stock, how='inner', on=['stock_code', 'date'])
    df = df[['order_book_id', 'stock_code', 'date', 'iv', 'RV']]
    df['IVdelta_mid_22D'] = df['iv'] - df['RV']
    df = df.dropna()
    return df

# 获取5D隐波差数据
def get_IV_RV_5D():
    df = IV_RV_5D()[['date', 'order_book_id', 'IVdelta_mid_5D']]
    return df

# 获取22D隐波差数据
def get_IV_RV_22D():
    df = IV_RV_22D()[['date', 'order_book_id', 'IVdelta_mid_22D']]
    return df

# 获取5D隐波差中位数
def get_IV_RV_5D_median():
    df = get_IV_RV_5D()
    df_IV_RV_median = df.groupby('date')['IVdelta_mid_5D'].median()
    df_IV_RV_median = df_IV_RV_median.reset_index().sort_values('date')
    return df_IV_RV_median

# 获取22D隐波差中位数
def get_IV_RV_22D_median():
    df = get_IV_RV_22D()
    df_IV_RV_median = df.groupby('date')['IVdelta_mid_22D'].median()
    df_IV_RV_median = df_IV_RV_median.reset_index().sort_values('date')
    return df_IV_RV_median

# 数据存档
def IVdelta_mid_5D_save():
    excel_file_path = 'IVdelta_mid_5D.csv'
    #if os.path.exists('原始数据/' + excel_file_path):
    #    df = pd.read_csv('原始数据/' + excel_file_path)
    #    return df
    #else:
    save_CSV(get_IV_RV_5D_median(), 'get_data/原始数据/' + excel_file_path.split('.')[0])
    df = pd.read_csv('get_data/原始数据/' + excel_file_path)
    return df

# 数据存档
def IVdelta_mid_22D_save():
    excel_file_path = 'IVdelta_mid_22D.csv'
    #if os.path.exists('原始数据/' + excel_file_path):
    #    df = pd.read_csv('原始数据/' + excel_file_path)
    #    return df
    #else:
    save_CSV(get_IV_RV_22D_median(), 'get_data/原始数据/' + excel_file_path.split('.')[0])
    df = pd.read_csv('get_data/原始数据/' + excel_file_path)
    return df

